"""Checkpointing support."""

import cPickle as pickle
import errno
import os
import os.path
import random
# pylint: disable-msg=W0402
import string
# pylint: enable-msg=W0402

_DEBUG = False

class CheckpointFile(object):
	# pylint: disable-msg=W0511
	"""A Checkpoint File.

	CheckpointFiles provide atomically-updated write-only file-like objects
	for maintaining checkpoint files. The idea of a checkpoint file is that
	it is always internally consistent -- updates to it happen atomically. On
	Unix, the way this works is that data is written to a temp file on the
	checkpoint file's file system, and then the temp file is rename()d to the
	checkpoint file name.

	CheckpointFile does not provide any facility whatsoever for reading the
	contents of an existing checkpoint file -- the file builtin in read-only
	mode is sufficient.

	The file object interface is followed by CheckpointFile but there are
	some additional semantics.

	When a CheckpointFile is created, it provides a handle to an empty file.
	Reading and writing to that file will behave as expected. At any time,
	commit() can be called. Once commit() is called, all data is flushed
	to disk and the checkpoint file is atomically updated. When this has
	succeeded, CheckpointFile once again becomes a handle to an empty file.

	Once commit() has been called, any read only references to the checkpoint
	file are immediately invalidated. It is the programmer's responsibility
	to account for this.

	A concrete example should make this clear.

	Let's create a CheckpointFile for /var/run/mydaemon.chk.

	chkf = CheckpointFile('/var/run/mydaemon.chk').

	chkf now reads and writes to a randomly-named temp file, say,
	/var/run/mydaemon.chk.XXXX

	chkf.write(somedata)

	/var/run/mydaemon.chk.XXXX now contains whatever "somedata" was.

	chkf.write(moredata)

	/var/run/mydaemon.chk.XXXX now contains whatever "somedata" was,
	concatenated with whatever "moredata" was.

	chkf.commit()

	/var/run/mydaemon.chk.XXXX is renamed to /var/run/mydaemon.chk.
	chkf now reads and writes to the beginning of the empty file
	/var/run/mydaemon.chk.YYYY.
	"""

	# pylint: enable-msg=W0511

	__masquerade_methods = [ 'read', 'write', 'readline', 'readlines',
		'writelines', 'seek', 'tell', 'truncate' ]

	def __init__(self, name):
		"""Create a checkpoint file that will commit to the given name."""
		self.rng = random.SystemRandom()
		self.dir = os.path.dirname(name)
		self.basename = os.path.basename(os.path.basename(name))

		self.name = name

		self.tmp_filename = ""
		self.tmp_file = None

		self.__open_tmp_file()

	def __open_tmp_file(self):
		"""Open the temp file that backs a CheckpointFile."""
		# The only safe way to open a temp file is to choose a random
		# name, and exclusively creat-open it. We loop until success.
		tmp_file_open = False
		while not tmp_file_open:
			try:
				# Attempt opening a randomly-named dotfile
				tmp_name = os.path.join(self.dir,
					'.' + self.basename + self.__gen_random_string())
				fdesc = os.open(tmp_name, os.O_RDWR | os.O_CREAT | os.O_EXCL)

				# We succeeded
				self.tmp_filename = tmp_name
				self.tmp_file = os.fdopen(fdesc, 'rb+')

				tmp_file_open = True

			except OSError, exc:
				# Ignore name collisions
				if exc.errno is not errno.EEXIST:
					raise

		# We masquerade as a file-like object by delegating the common
		# __builtin__.file methods to the temp file object.
		for method_name in CheckpointFile.__masquerade_methods:
			setattr(self, method_name, getattr(self.tmp_file, method_name))

	def __gen_random_string(self, length=20):
		"""Generate a random string of the given length from the set of
			alphanumerics."""

		return ''.join(self.rng.sample(string.letters + string.digits, length))

	def commit(self):
		"""Write all pending changes to the checkpoint file.

		After a commit, the caller will be left with a file handle to the
		beginning of a truncated file."""
		
		self.tmp_file.flush()
		self.tmp_file.close()
		os.rename(self.tmp_filename, self.name)
		self.__open_tmp_file()

	def close(self):
		"""Close this checkpoint file.

		NB: close() does not imply commit()."""

		self.tmp_file.close()
		os.unlink(self.tmp_filename)

def dump(obj, checkpoint_file, protocol=None):
	"""Pickle the given object to the provided checkpoint file."""
	checkpoint_file.seek(0)
	checkpoint_file.truncate()

	if protocol is None:
		pickle.dump(obj, checkpoint_file)
	else:
		pickle.dump(obj, checkpoint_file, protocol)

	checkpoint_file.commit()

class CheckpointCollection(object):
	"""CheckpointCollection is a mixin class that allows collections,
	like dictionaries, lists, and so on, to provide a method for commit()
	themselves to a checkpoint file.

	Comparitors for CheckpointCollections are inherited from the collection
	used.
	"""

	def __init__(self, checkpoint_file):
		"""Initialize this CheckpointCollection from the provided
		checkpoint file."""
		self.checkpoint_file = checkpoint_file

	def set_checkpoint_file(self, checkpoint_file):
		"""Associate a new CheckpointFile with this collection."""
		self.checkpoint_file = checkpoint_file

	def commit(self):
		"""Pickle this collection to its CheckpointFile."""
		if self.checkpoint_file:
			chkfile = self.checkpoint_file
			self.checkpoint_file = None
			dump(self, chkfile)
			self.checkpoint_file = chkfile
		else:
			raise TypeError("Can not commit() an unpickled %s "
						"without first set_checkpoint_file()ing it."
						% type(self).__name__
					)

	def close(self):
		"""Close the CheckpointCollection and its underlying file.
		The same semantics apply to this as to CheckpointFile.close()."""
		self.checkpoint_file.close()

class CheckpointDictionary(CheckpointCollection, dict):
	"""A CheckpointCollection of a dict."""
	def __init__(self, checkpoint_file, *args, **kwargs):
		"""Create the CheckpointDictionary.

		CheckpointCollection consumes the checkpoint_file argument.
		All other arguments are passed directly to the dictionary
		constructor."""

		CheckpointCollection.__init__(self, checkpoint_file)
		dict.__init__(self, *args, **kwargs)

class CheckpointList(CheckpointCollection, list):
	"""A CheckpointCollection of a list."""
	def __init__(self, checkpoint_file, *args, **kwargs):
		"""Craete the CheckpointList.

		CheckpointCollection consumes the checkpoint_file argument.
		All other arguments are passed directly to the list constructor."""
		CheckpointCollection.__init__(self, checkpoint_file)
		list.__init__(self, *args, **kwargs)
